import argparse
import os
import sys
import uuid
import csv
import time
from threading import Barrier
from locust import HttpUser, task, between, events, LoadTestShape
from locust.exception import StopUser
from config.all_paths import get_root_path
from utils.file_handling.excel_handling import ExcelUtil
from utils.interface_auth.auth_tool import Authentication
from utils.performance_tool.config_loader import load_config
from utils.performance_tool.load_shapes import SHAPE_MAP
from utils.request_tool.send_request import RequestSender
from utils.time_tool.time_control import parse_duration

# =================== 命令行参数 ===================
def parse_custom_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', required=True, help='环境名')
    parser.add_argument('--role', required=True, help='测试角色')
    parser.add_argument('--profile_name', required=True, help='压测方案名')
    args, remaining = parser.parse_known_args()
    sys.argv = sys.argv[:1] + remaining
    return args

args = parse_custom_args()
ENV, ROLE, PROFILE_NAME = args.env, args.role, args.profile_name

# =================== 加载配置 ===================
config = load_config(env=ENV, role=ROLE)
env_config = config['env']
profile_config = config['profiles'].get(PROFILE_NAME, {})

# =================== 读取测试数据 ===================
excel_path = os.path.join(get_root_path(), '../data/case_excel/wutong/uat/医生账号相关接口.xlsx')
testData = ExcelUtil(excel_path).read_excel()

# =================== 认证 ===================
class BaseUser(HttpUser):
    wait_time = between(0.5, 2)
    abstract = True

    def on_start(self):
        auth = Authentication(
            env_config.get('username'),
            env_config.get('password'),
            env_config.get('url')
        )
        auth_method_map = {
            "test_wt": "auth_wutong",
            "test_th": "auth_th",
            "test_inter": "auth_inter"
        }
        method = auth_method_map.get(ENV, "auth_wutong")
        self.s = getattr(auth, method)()
        self.user_id = str(uuid.uuid4())  # 为每个用户生成唯一ID

# =================== 用户行为类 ===================
class BaselineUser(BaseUser):
    @task
    def task_run(self):
        print(f"[Baseline] 用户 {self.user_id} 执行请求")
        RequestSender().send_request(self.s, testData[0])

class StressUser(BaseUser):
    @task(3)
    def task_run(self):
        print(f"[Stress] 用户 {self.user_id} 执行请求")
        RequestSender().send_request(self.s, testData[0])

class WaveUser(BaseUser):
    @task
    def task_run(self):
        print(f"[Wave] 用户 {self.user_id} 执行请求")
        RequestSender().send_request(self.s, testData[0])

# 集合点 Barrier（支持并发测试）
barrier = Barrier(parties=profile_config.get("users", 1))

class ConcurrencyUser(BaseUser):
    @task
    def task_run(self):
        try:
            print(f"[Concurrency] 用户 {self.user_id} 等待集合点...")
            barrier.wait()
            print(f"[Concurrency] 用户 {self.user_id} 发起请求！")
            RequestSender().send_request(self.s, testData[0])
        except Exception as e:
            print(f"[并发异常] {e}")
            raise StopUser()

# =================== 动态类和形状类 ===================
class DynamicTestClass(eval(profile_config.get("user_class", "BaselineUser"))):
    pass

class CustomLoadShape(SHAPE_MAP[profile_config.get("load_shape", "constant")]):
    shape_users = profile_config.get("users", 1)
    spawn_rate = profile_config.get("spawn_rate", 1)
    shape_duration = parse_duration(profile_config.get("duration", "1m"))
    shape_stages = profile_config.get("stages", [])

# 暴露给Locust的类
class TestScenario(DynamicTestClass):
    pass

def get_shape_class():
    return CustomLoadShape

# =================== 压测统计与报告 ===================
results = []

@events.request.add_listener
def on_request(request_type, name, response_time, response_length, **kwargs):
    results.append({
        "type": request_type,
        "name": name,
        "response_time": response_time,
        "response_length": response_length,
        "timestamp": time.time()
    })

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("\n=== 测试结束，统计结果如下 ===")
    total_requests = environment.stats.total.num_requests
    fail_ratio = environment.stats.total.fail_ratio
    avg_response_time = environment.stats.total.avg_response_time
    tps = total_requests / environment.stats.total.total_rps if environment.stats.total.total_rps else 0

    print(f"请求总数: {total_requests}")
    print(f"失败率: {fail_ratio:.2%}")
    print(f"平均响应时间: {avg_response_time:.2f}ms")
    print(f"TPS（每秒事务数）: {tps:.2f}")

    # 导出结果 CSV
    report_path = os.path.join(get_root_path(), f"locust_report_{PROFILE_NAME}.csv")
    with open(report_path, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.DictWriter(file, fieldnames=results[0].keys())
        writer.writeheader()
        writer.writerows(results)
    print(f"报告已生成: {report_path}")
